import json
import os
import shutil
from glob import glob
from pathlib import Path

from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

from option import parse_args
from preprocessing import face_detector
from preprocessing.face_detector import VideoDataset


def temp_func(x):
    return x


def process_videos(videos, root_dir, detector_cls: str):
    detector = face_detector.__dict__[detector_cls](device="cuda:0")
    dataset = VideoDataset(videos)
    loader = DataLoader(dataset, shuffle=False, num_workers=1, batch_size=1, collate_fn=temp_func)
    for item in tqdm(loader):
        video, indices, frames = item[0]
        video_path_split = os.path.splitext(os.path.basename(video))
        video_id = video_path_split[0]
        part_id = Path(video).parent.name
        out_dir = os.path.join(root_dir, "boxes", part_id)
        output_json = os.path.join(out_dir, "{}.json".format(video_id))
        if os.path.exists(output_json):
            continue
        batches = [frames[i:i + detector.batch_size] for i in range(0, len(frames), detector.batch_size)]
        result = {}
        for j, frames in enumerate(batches):
            result.update({int(j * detector.batch_size) + i: b for i, b in zip(indices, detector.detect_faces(frames))})
        os.makedirs(out_dir, exist_ok=True)
        with open(os.path.join(out_dir, "{}.json".format(video_id)), "w") as f:
            json.dump(result, f)


def move_boxes(videos, root_dir, detector_type):
    # dataset = VideoDataset(videos)
    # loader = DataLoader(dataset, shuffle=False, num_workers=1, batch_size=1, collate_fn=temp_func)
    for item in tqdm(videos):
        video = item
        video_path_split = os.path.splitext(os.path.basename(video))
        video_id = video_path_split[0]
        part_id = Path(video).parent.name

        boxes_dir = os.path.join(root_dir, "boxes")
        video_boxes_json = os.path.join(boxes_dir, "{}.json".format(video_id))
        boxes_part_dir = os.path.join(boxes_dir, part_id)
        if not os.path.exists(boxes_part_dir):
            os.makedirs(boxes_part_dir, exist_ok=True)
        part_video_boxes_json = os.path.join(boxes_part_dir, "{}.json".format(video_id))

        if os.path.exists(video_boxes_json):
            shutil.move(video_boxes_json, part_video_boxes_json)


def get_original_video_paths(root_dir, basename=False):
    video_dir = os.path.join(root_dir, 'videos')
    originals = set()
    originals_v = set()
    for json_path in glob(os.path.join(video_dir, "*/metadata.json")):
        parent_dir = Path(json_path).parent
        with open(json_path, "r") as f:
            metadata = json.load(f)
        for k, v in metadata.items():
            if v["label"] == "REAL":
                original = k
                originals_v.add(original)
                originals.add(os.path.join(parent_dir, original))
    originals = list(originals)
    originals_v = list(originals_v)
    return originals_v if basename else originals


def main():
    args = parse_args()
    original_videos = get_original_video_paths(args.root_dir)
    print(len(original_videos))
    move_boxes(original_videos, args.root_dir, args.detector_type)
    # process_videos(original_videos, args.root_dir, args.detector_type)


if __name__ == "__main__":
    main()
